Summary

Wave 1 of the staged post-load hooks rollout. The staged-hook contract landed in #14770, and #14878 has now merged, so this PR is a single Wave 1 commit on top of main.

This change migrates alias-only model hooks from post_load_weights() to setup_aliases() and cuts the GMS read-only (RO) load path over from the old meta-tensor workaround to the staged-hook protocol.

What Changed

• Alias migration for 7 top-level model classes: modeling_llama (LlamaForCausalLM, Llama4ForConditionalGeneration), modeling_deepseekv3, modeling_glm, modeling_exaone_moe, modeling_qwen3_moe, modeling_qwen3_next, and modeling_gpt_oss. Their alias-only post_load_weights() bodies move verbatim into setup_aliases(), while standard load paths continue through the base post_load_weights() orchestrator.
• GMS RO ordering now runs staged hooks around zero-copy materialization:

post_load_apply
  -> _setup_aliases(model)                  # recursive alias walk
  -> _check_gms_source_identity(gms_backend) # STRICT pre-materialize gate from #14878
  -> materialize_module(model)
  -> _walk_cache_state(model)
  -> post_load_publish

• GMS docs now describe alias setup before materialization and derived-state refresh after real tensors are bound.
• Full reload now resets existing _weights_transformed flags before rebinding fresh weights, while partial reload keeps existing transform guards intact for untouched modules.
• Tests cover the staged walkers, reload reset, and GMS RO ordering in test_model_loader_gms.py and test_model_loader_mx.py.

Dependency / prerequisite stack

This PR is Wave 1 in the staged post-load hooks rollout. The foundation PRs #14770 and #14878 are already merged. The wave PRs should merge in sequence; after each upstream wave lands, rebase the next wave onto main so review and CI focus on that wave's delta.

Arrows point from prerequisite to dependent. PR numbers in graph nodes are clickable.

graph TD
    PR14770["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/14770'>#14770</a>: staged-hook contract (merged)"]
    PR14878["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/14878'>#14878</a>: GMS SourceIdentity gate (merged)"]
    PR15014["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15014'>#15014</a>: Wave 1 aliases + GMS RO load (this PR, open)"]
    PR15288["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15288'>#15288</a>: Wave 2 Linear/Attention transforms (draft)"]
    PR15386["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15386'>#15386</a>: Wave 3 MoE/Mamba staged hooks (draft)"]
    PR15387["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15387'>#15387</a>: Wave 4 MX receiver cutover (draft)"]
    PR15432["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15432'>#15432</a>: Wave 5 MX publisher + Llama receiver (draft)"]
    VERIFY["post-migration verification / demo (planned)"]

    PR14770 -->|satisfied| PR15014
    PR14878 -->|satisfied| PR15014
    PR15014 -->|blocking| PR15288
    PR15288 -->|blocking| PR15386
    PR15386 -->|blocking| PR15387
    PR15387 -->|blocking| PR15432
    PR15432 -.->|planned| VERIFY

    classDef merged fill:#dcfce7,stroke:#16a34a,color:#14532d;
    classDef inflight fill:#dbeafe,stroke:#2563eb,color:#1e3a8a;
    classDef draft fill:#ffedd5,stroke:#f97316,color:#7c2d12;
    classDef current fill:#ede9fe,stroke:#7c3aed,color:#3b0764,stroke-width:3px;
    classDef downstream fill:#f3f4f6,stroke:#6b7280,color:#374151,stroke-dasharray:5 5;
    linkStyle 0,1 stroke:#16a34a,stroke-width:2px;
    linkStyle 2,3,4,5 stroke:#ea580c,stroke-width:3px;
    linkStyle 6 stroke:#6b7280,stroke-width:2px,stroke-dasharray:5 5;

    class PR14770,PR14878 merged;
    class PR15288,PR15386,PR15387,PR15432 draft;
    class PR15014 current;
    class VERIFY downstream;

Immediate merge dependency for this PR: none beyond the already-merged #14770 and #14878 foundation PRs; downstream waves remain stacked on this branch until Wave 1 lands.

Test Plan

• pytest tests/unittest/_torch/pyexecutor/test_model_loader_gms.py tests/unittest/_torch/pyexecutor/test_model_loader_mx.py
• Full L0 CI with /bot run --disable-fail-fast before review.
• Local checks completed: git diff --check, py_compile on changed Python files, and pre-commit. Local pytest was not available in this macOS shell.

Next Steps

• Wave 2: migrate Linear/Attention transforms into transform_weights() with _weights_transformed guards.
• Wave 3: migrate MoE and Mamba transforms.
• Wave 4: MX publish-after-transform flip, receiver cutover, and per-model allow-list.

PR Checklist

• PR description clearly explains what and why.
• Follows TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No public API changes.
• No new dependencies.

Summary by CodeRabbit

• Refactor
  • Restructured internal model initialization and layer alias resolution for improved stability during weight loading.
  • Optimized GPU memory management during model setup and state caching.
  • Updated model loading sequence for read-only GPU memory paths to ensure correct operation ordering.